﻿from memory import GlobalMemory
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple
from enum import Enum
import numpy as np
from numpy import ndarray, uint8, uint16, uint32, uint64, float32, float64
import threading
from datetime import datetime


class MESI(Enum):
    Invalid = 'I'
    Exclusive = 'E'
    Shared = 'S'
    Modified = 'M'


@dataclass
class CacheLine:
    tag: int = 0
    state: MESI = MESI.Invalid
    data: ndarray[uint8] = field(default_factory=lambda: np.zeros(
        shape=64, dtype=uint8))


@dataclass
class LRUCache:
    def __init__(self, size: int, memory: GlobalMemory):
        self.__history: Dict[int, int] = {}
        self.__data: Dict[int, CacheLine] = {}
        self.__size: int = size
        self.__memory = memory

    def __getitem__(self, addr: int) -> CacheLine:
        if addr not in self.__data.keys():
            if len(self.__data.keys()) == self.__size:
                id = min(self.__history.keys(), key=self.__history.get)
                self.__history.pop(id)
                if self.__data[id].state == MESI.Modified:
                    self.__memory.put(
                        self.__data[id].tag, self.__data[id].data)
                self.__data.pop(id)
            self.__data[addr] = CacheLine(
                addr >> 12, MESI.Exclusive, self.__memory.get(addr))
        self.__history[addr] = datetime.now().timestamp()
        return self.__data[addr]

    def __setitem__(self, addr: int, data: CacheLine):
        if addr in self.__data.keys():
            if self.__data[addr].state == MESI.Modified:
                self.__memory.put(
                    self.__data[addr].tag, self.__data[addr].data)
        self.__data[addr] = data
        self.__history[addr] = datetime.now().timestamp()

    def __delitem__(self, addr: int):
        if addr in self.__data.keys():
            if self.__data[addr].state == MESI.Modified:
                self.__memory.put(
                    self.__data[addr].tag, self.__data[addr].data)
            self.__data.pop(addr)
            self.__history.pop(addr)


@dataclass
class Cache:
    def __init__(self, memory: GlobalMemory):
        self.__l1i: Dict[int, CacheLine] = {}
        self.__l1d: Dict[int, LRUCache] = {}
        self.__l2: Dict[int, LRUCache] = {}
        self.__memory = memory

    def read(self, addr: int, size: int, is_ins: bool) -> ndarray[uint8]:
        tag_id = addr >> 6
        if is_ins:
            idx = tag_id & 0x3ff
            if idx not in self.__l1i.keys():
                self.__l1i[idx] = CacheLine(
                    tag_id, MESI.Invalid, np.zeros(shape=64, dtype=uint8))
            line = self.__l1i[idx]
            if line.state == MESI.Invalid:
                idx = tag_id & 0x3fff
                if idx not in self.__l2.keys():
                    self.__l2[idx] = LRUCache(2, self.__memory)
                line = self.__l2[idx][tag_id]
                self.__l1i[tag_id & 0x3ff] = line
                return line.data[addr & 0x3c:(addr & 0x3c)+4]
            else:
                return line.data[addr & 0x3c:(addr & 0x3c)+4]
        else:
            base_idx = addr & (0x40-size)
            idx = tag_id & 0x3ff
            if idx not in self.__l1d.keys():
                self.__l1d[idx] = LRUCache(4, self.__memory)
            line = self.__l1d[idx][tag_id]
            if line.state == MESI.Invalid:
                idx = tag_id & 0x3fff
                if idx not in self.__l2.keys():
                    self.__l2[idx] = LRUCache(2, self.__memory)
                line = self.__l2[idx][tag_id]
                self.__l1d[tag_id & 0x3ff][tag_id] = line
                return line.data[base_idx:addr & base_idx+size]
            else:
                return line.data[base_idx:base_idx+size]

    def write(self, addr: int, size: int, data: ndarray[uint8]):
        tag_id = addr >> 6
        idx = tag_id & 0x3ff
        base_idx = addr & (0x40-size)
        line = self.__l1d[idx][tag_id]
        if line.state == MESI.Invalid:
            idx = tag_id & 0x3fff
            line = self.__l2[idx][tag_id]
            if len(data) != size:
                raise ValueError(f"data size not match")
            line.state = MESI.Modified
            line.data[base_idx:base_idx+size] = data
            self.__l1d[tag_id & 0x3ff][tag_id] = line
        else:
            line.state = MESI.Modified
            line.data[base_idx:base_idx+size] = data
            self.__l1d[tag_id & 0x3ff][tag_id] = line
